package com.ls.aicoderbackend.ai;

import com.github.benmanes.caffeine.cache.Cache;
import com.github.benmanes.caffeine.cache.Caffeine;
import com.ls.aicoderbackend.ai.guardrail.PromptSafetyInputGuardrail;
import com.ls.aicoderbackend.ai.guardrail.RetryOutputGuardrail;
import com.ls.aicoderbackend.ai.tools.ToolManager;
import com.ls.aicoderbackend.exception.BusinessException;
import com.ls.aicoderbackend.exception.ErrorCode;
import com.ls.aicoderbackend.model.enums.CodeGenTypeEnum;
import com.ls.aicoderbackend.service.ChatHistoryService;
import com.ls.aicoderbackend.utils.SpringContextUtil;
import dev.langchain4j.community.store.memory.chat.redis.RedisChatMemoryStore;
import dev.langchain4j.data.message.ToolExecutionResultMessage;
import dev.langchain4j.memory.chat.MessageWindowChatMemory;
import dev.langchain4j.model.chat.ChatModel;
import dev.langchain4j.model.chat.StreamingChatModel;
import dev.langchain4j.service.AiServices;
import jakarta.annotation.Resource;
import lombok.extern.slf4j.Slf4j;
import org.springframework.context.annotation.Bean;
import org.springframework.context.annotation.Configuration;

import java.time.Duration;

@Configuration
@Slf4j
public class AICodeGeneratorServiceFactory {

    @Resource(name = "openAiChatModel")
    private ChatModel chatModel;

    @Resource
    private RedisChatMemoryStore redisChatMemoryStore;

    @Resource
    private ChatHistoryService chatHistoryService;

    @Resource
    private ToolManager toolManager;


    /**
     * 缓存 AiCodeGeneratorService 实例
     */
    private final Cache<String, AiCodeGeneratorService> serviceCache = Caffeine.newBuilder()
            .maximumSize(1000)
            .expireAfterWrite(Duration.ofMinutes(30))
            .expireAfterAccess(Duration.ofMinutes(10))
            .removalListener((key, val, cause) -> {
                log.debug("AiService 实例被移除，appId：{}， 原因：{}", key, cause);
            }).build();


    private AiCodeGeneratorService createAiCodeGeneratorService(Long appId, CodeGenTypeEnum codeGenTypeEnum) {
        MessageWindowChatMemory chatMemory = MessageWindowChatMemory.builder()
                .id(appId)
                .chatMemoryStore(redisChatMemoryStore)
                .maxMessages(20)
                .build();
        // 从缓存中获取对应appId的对话记录
        chatHistoryService.loadChatHistoryMemoryToRedis(appId, chatMemory, 20);
        // 返回AiService服务，通过动态代理获得
        return switch (codeGenTypeEnum) {
            case HTML, MULTI_FILE -> {
                // 使用多例模式注入Bean解决并发问题
                StreamingChatModel openAiStreamingChatModel = SpringContextUtil.getBean("streamingChatModelPrototype", StreamingChatModel.class);
                yield AiServices.builder(AiCodeGeneratorService.class)
                        .chatModel(chatModel)
                        .streamingChatModel(openAiStreamingChatModel)
                        .chatMemory(chatMemory)
                        .inputGuardrails(new PromptSafetyInputGuardrail())
//                      .outputGuardrails(new RetryOutputGuardrail()) 使用输出护轨会导致无法流式输出
                        .build();
            }
            case VUE_PROJECT -> {
                // 使用多例模式注入Bean解决并发问题
                StreamingChatModel reasoningStreamingChatModel = SpringContextUtil.getBean("reasoningStreamingChatModelPrototype", StreamingChatModel.class);
                yield AiServices.builder(AiCodeGeneratorService.class)
                        .chatModel(chatModel)
                        .streamingChatModel(reasoningStreamingChatModel)
                        .chatMemoryProvider(memoryId -> chatMemory)
                        .tools(toolManager.getAllTools())
                        .hallucinatedToolNameStrategy(toolExecutionRequest ->
                                ToolExecutionResultMessage.from(toolExecutionRequest,
                                        "Error: there is no tool call" + toolExecutionRequest.name()
                                )
                        )
                        .inputGuardrails(new PromptSafetyInputGuardrail())
                        .maxSequentialToolsInvocations(20) // 最多允许ai调用20次工具
//                      .outputGuardrails(new RetryOutputGuardrail()) 使用输出护轨会导致无法流式输出
                        .build();
            }
            default -> throw new BusinessException(ErrorCode.SYSTEM_ERROR, "不支持的代码生成类型");
        };
    }


    public AiCodeGeneratorService getAiCodeGeneratorService(Long appId) {
        return getAiCodeGeneratorService(appId, CodeGenTypeEnum.HTML);
    }

    /**
     * 根据appId和代码生成类型获取对应的AiCodeGeneratorService实例
     * @param appId
     * @return
     */
    public AiCodeGeneratorService getAiCodeGeneratorService(Long appId, CodeGenTypeEnum codeGenTypeEnum) {
        String cacheKey = buildCacheKey(appId, codeGenTypeEnum);
        return serviceCache.get(cacheKey, key -> createAiCodeGeneratorService(appId, codeGenTypeEnum));
    }

    // 构建缓存Key
    private String buildCacheKey(Long appId, CodeGenTypeEnum codeGenTypeEnum) {
        return appId + "_" + codeGenTypeEnum.getValue();
    }

    @Bean
    public AiCodeGeneratorService createAICodeGeneratorService() {
        return getAiCodeGeneratorService(0L);
    }
}
